#!/usr/bin/env python
# -*- encoding: UTF-8 -*-
from __future__ import print_function
__doc__='''
松协调: 初始位姿我摆出来,则可通过相对位姿确定个机械臂的角度,

左手上拿个棋盘格，右手上放一个标记点，

左手运动时，保持右手上标记点画圆，

实际上就是右手标记点沿给定轨迹运动
'''
import sql
import cv2
import numpy as np
import os
import csv
import robot
import common
from datetime import datetime
basedir = os.path.abspath(os.path.dirname(__file__))
state_dir = os.path.join(basedir, "data", "State")

class Cooperation(object):
    def __init__(self,nao):
        assert isinstance(nao,sql.Nao)
        self.nao = nao
        self.use_cam_p = False
        self.use_cam_pp = True
        if self.use_cam_p:
            self.cap = cv2.VideoCapture(0) # 外加相机 +
        if self.use_cam_pp:
            self.capp = cv2.VideoCapture(0) # 外加相机 ++
        self.w = 9 # 棋盘格宽
        self.h = 4 # 棋盘格高
        self.unit = 15 # 棋盘格单位长度 mm
        self.dt = 0.01 # 仿真精度
        self.animation = False # 手动调整姿态标记

    def get_image(self):
        self._top_image = self.nao.top_cam.get_image()
        self._bottom_image = self.nao.bottom_cam.get_image()
        if self.use_cam_p:
            self._camp_image = self.cap.read()[1]
        if self.use_cam_pp:
            self._campp_image = self.capp.read()[1]
        return True

    def show(self):
        _top = self._top_image.copy()
        _bottom = self._bottom_image.copy()
        common.drawChessboard(_bottom,w=self.w,h=self.h)
        cv2.imshow("top",_top)
        cv2.imshow("bottom",_bottom)
        if self.use_cam_p:
            cv2.imshow("cam_p",self._camp_image)
        if self.use_cam_pp:
            cv2.imshow("cam_pp",self._campp_image)

    def save(self,name=None,now=None):
        if now is None:
            now = datetime.now()
        if name is None:
            name = "songxietiao"
        else:
            name = "songxietiao"+name
        self.nao.states.append(sql.State(name,self.nao,
                                         w=self.w,h=self.h,
                                         unit=self.unit,now=now))
        time_name = str(now).replace(':','.') # : 不能用作 win32 的文件名
        if self.use_cam_p:
            img_name = " ".join([time_name,name,self.nao.name,'camp'+'.jpg'])
            cv2.imwrite(os.path.join(state_dir,img_name),self._camp_image)
        if self.use_cam_pp:
            img_name = " ".join([time_name,name,self.nao.name,'campp'+'.jpg'])
            cv2.imwrite(os.path.join(state_dir,img_name),self._campp_image)

    def random(self,q0,delta=0.1,left=True):
        '''
        随机在 q0 附近，各关节角为delta附近运动一次，
        用来求图像雅克比的初始值，和估计标记点的位姿等
        :param q0:  初始值
        :param delta: 每个关节的变化范围
        :return:
        '''
        if isinstance(q0,list):
            q0 = np.array(q0)
        dq = np.random.random(q0.shape)*delta
        dq = dq - delta/2.0
        q = q0+dq
        q[-1] = 0 # 默认手抓不张开
        if left:
            self.nao._setAngles("LArm",q.tolist())
        else:
            self.nao._setAngles("RArm",q.tolist())

    def commit(self):
        ss = sql.Session()
        ss.add(self.nao)
        ss.commit()


def connect(nao_name):
    nao = sql.Nao.get_by_name(nao_name)
    nao.connect = True
    ss = sql.Session()
    ss.add(nao)
    ss.commit()
    ss.close()

def disconnect(nao_name):
    nao = sql.Nao.get_by_name(nao_name)
    nao.connect = False
    ss = sql.Session()
    ss.add(nao)
    ss.commit()
    ss.close()

import time
import thread

def animation_func(coop,touchProxy,motionProxy):
    assert isinstance(coop,Cooperation)
    def switch_stiffnesses(name):
        stifffs = motionProxy.getStiffnesses(name)
        if isinstance(stifffs,list):
            stifffs = np.linalg.norm(stifffs)
        if stifffs < 0.5:
            motionProxy.setStiffnesses(name,0.8)
        else:
            motionProxy.setStiffnesses(name,0.0)

    while coop.animation:
        print("coop.animation on")
        time.sleep(1.0)
        for state in touchProxy.getStatus():
            if state[1]:
                if state[0] == 'Head':
                    switch_stiffnesses('Head')
                    break # 按下了头部的传感器
                elif state[0] == 'LHand/Touch/Back':
                    switch_stiffnesses('LArm')
                    break
                elif state[0] == 'RHand/Touch/Back':
                    switch_stiffnesses('RArm')
                    break
                elif state[0] in ['LFoot/Bumper/Left','LFoot/Bumper/Right']:
                    switch_stiffnesses('LLeg')
                    break
                elif state[0] in ['RFoot/Bumper/Left','RFoot/Bumper/Right']:
                    switch_stiffnesses('RLeg')
                    break

ch_list = []
g_dict = {}
def animation(string,coop,ch):
    string+="animation: o[n]/o[f]f"
    if ch == 'n' and not coop.animation:
        coop.animation = True
        touchProxy = coop.nao.touchProxy
        motionProxy = coop.nao.motionProxy
        thread.start_new_thread(animation_func,(coop,touchProxy,motionProxy))
    elif ch == 'f' and coop.animation:
        coop.animation = False
    return string

def pase_command(ch,coop):
    '''命令解析'''
    isinstance(coop,Cooperation)
    g_dict["save_name"] = "xtag"
    g_dict["random_delta"] = 0.1
    ch = chr(ch).lower() # 字符化，不区分大小写
    string = "[b]ack:" # 输出字符串
    print(ch_list)
    if len(ch_list) > 0:
        if ch_list[0] in ['n']:
            string += "%s:" % coop.nao.name
        elif ch_list[0] == 'c':
            string += "Cooperation:"
        elif ch_list[0] == 'f':
            string+="Config:"
        # 退回
        if ch == 'b':
            ch_list.pop()
    # 解析并执行命令
    if len(ch_list) == 0:
        print("choose control [n]ao/[C]ooperation/con[f]ig")
        if ch in ['n','c','f']:
            ch_list.append(ch)
            return None
    elif len(ch_list) == 1:
        if ch_list[0] == 'n':
            string += "[m]otion/"
            string += "[p]osture/"
            string += "[a]nimation/"
            string += "[d]isableAutonomousLife/"
            if ch in ['m','a']:
                ch_list.append(ch)
            elif ch == 'd':
                coop.nao._get_proxy("ALAutonomousLife").setState("disabled")
        elif ch_list[0] == 'c':
            string += "[s]ave/[c]ontrol/[a]nimation"
            if ch in ['c','a']:
                ch_list.append(ch)
            elif ch == 's':
                coop.save(g_dict["save_name"])
                coop.commit()
        elif ch_list[0] == 'f':
            string+=" [n]ame/[q]0/[d]elta"
            if ch in ['n','d']:
                ch_list.append(ch)
            if ch == 'q':
                g_dict["left_q0"] = coop.nao._getAngles("LArm")
                g_dict["right_q0"] = coop.nao._getAngles("RArm")
                coop.save(g_dict["save_name"]+"_q0")
                coop.commit()
    elif len(ch_list) == 2:
        if ch_list[0] == 'n':
            if ch_list[1] == 'm':
                string +="motion: "
                string += "[w]akeUp/"
                string += "[r]est/"
                string += "[h]and/"
                if ch == 'w':
                    coop.nao.motionProxy.wakeUp()
                elif ch == 'r':
                    coop.nao.motionProxy.rest()
            elif ch_list[1] == 'p':
                string +="posture: "
                string += "[S]tand/"
                string += "Stand[I]nit/"
                string += "[C]rouch"
                if ch == 's':
                    coop.nao.postureProxy.goToPosture("Stand", 1.0)
                elif ch == 'i':
                    coop.nao.postureProxy.goToPosture("StandInit", 1.0)
                elif ch == 'c':
                    coop.nao.postureProxy.goToPosture("Crouch", 1.0)
            elif ch_list[1] == 'a':
                string = animation(string,coop,ch)
        elif ch_list[0] == 'c':
            if ch_list[1] == 'a':
                string = animation(string,coop,ch)
            elif ch_list[1] == 'c':
                string+="control: [s]ave/set[q]0/random[L]eft/random[R]ight"
                delta = g_dict["random_delta"]
                if ch == 's':
                    coop.save(g_dict["save_name"]+"_random")
                    coop.commit()
                elif ch == 'q':
                    g_dict["left_q0"] = coop.nao._getAngles("LArm")
                    g_dict["right_q0"] = coop.nao._getAngles("RArm")
                    coop.save(g_dict["save_name"]+"_q0")
                    coop.commit()
                elif ch == 'l':
                    if "left_q0" not in g_dict:
                        g_dict["left_q0"] = coop.nao._getAngles("LArm")
                    q0 = g_dict["left_q0"]
                    coop.random(q0=q0,delta=delta,left=True)
                elif ch == 'r':
                    if "right_q0" not in g_dict:
                        g_dict["right_q0"] = coop.nao._getAngles("RArm")
                    q0 = g_dict["right_q0"]
                    coop.random(q0=q0,delta=delta,left=False)
        elif ch_list[0] == 'f':
            if ch_list[1] == 'n':
                string+="change name [m]ark8/[x]tag/[c]hessboard"
                if ch == 'm':
                    g_dict["save_name"] = 'mark8'
                elif ch == 'x':
                    g_dict["save_name"] = 'xtag'
                elif ch == 'c':
                    g_dict["save_name"] = 'chessboard'
            elif ch_list[1] == 'd':
                string+= "delta=%f:" % g_dict["random_delta"]
                string+= "[a]dd/[m]inus"
                if ch == 'a':
                    g_dict["random_delta"] += 0.01
                elif ch == 'm':
                    g_dict["random_delta"] -= 0.01
    elif len(ch_list) == 3:
        if ch_list[0] == 'n':
            if ch_list[1] == 'm' and ch_list[2] == 'h':
                string+="Hand: "
                string+="openR[i]ght/closeRi[g]ht/openL[e]ft/closeLe[f]t"
                if ch == 'i':
                    coop.nao.motionProxy.openHand("RHand")
                elif ch == 'g':
                    coop.nao.motionProxy.closeHand("RHand")
                elif ch == 'e':
                    coop.nao.motionProxy.openHand("LHand")
                elif ch == 'f':
                    coop.nao.motionProxy.closeHand("LHand")
    print(string)

if __name__ == '__main__':
    nao_name = 'red'
    connect(nao_name)
    nao = sql.Nao.get_by_name(nao_name)
    coop = Cooperation(nao)
    while True:
        coop.get_image()
        coop.show()
        ch = cv2.waitKey(100) & 0xFF
        if ch == 27:
            break
        else:
            pase_command(ch,coop)
    disconnect(nao_name)


# def chaiji(nao_name,w=12,h=9,unit=20):
#     get_nao = lambda : sql.Nao.get_by_name(nao_name)
#     nao = get_nao()
#     autonomousLifeProxy = nao._get_proxy("ALAutonomousLife")
#     postureProxy = nao._get_proxy("ALRobotPosture")
#     step = np.deg2rad(3)
#     minHeadYaw, maxHeadYaw, minHeadPitch,maxHeadPitch = None,None,None,None
#     while True:
#         top_img = nao.top_cam.get_image()
#         bottom_img = nao.bottom_cam.get_image()
#         motionProxy = nao._get_proxy("ALMotion")
#         HeadYaw,HeadPitch = motionProxy.getAngles("Head",True)
#         Tb_o,Tc_o,position = nao.findTr(w,h,unit,top_img,bottom_img)
#         drawChessboard(top_img,w=12,h=9)
#         drawChessboard(bottom_img,w=12,h=9)
#         draw_str(top_img,(10,10),"HeadYaw:"+str(HeadYaw))
#         draw_str(top_img,(10,25),"HeadPitch:"+str(HeadPitch))
#         draw_str(top_img,(10,40),"minHeadYaw:"+str(minHeadYaw))
#         draw_str(top_img,(10,55),"maxHeadYaw:"+str(maxHeadYaw))
#         draw_str(top_img,(10,70),"minHeadPitch:"+str(minHeadPitch))
#         draw_str(top_img,(10,85),"maxHeadPitch:"+str(maxHeadPitch))
#         draw_str(bottom_img,(10,10),"Tb_o:")
#         if Tb_o is not None:
#             for x,p in enumerate(np.round(Tb_o,3).tolist()):
#                 draw_str(bottom_img,(10,25+x*15),str(p))
#         draw_str(bottom_img,(10,90),"Tc_o:")
#         if Tc_o is not None:
#             for x,p in enumerate(np.round(Tc_o,3).tolist()):
#                 draw_str(bottom_img,(10,105+x*15),str(p))
#             if position == "top":
#                 used_cam = nao.top_cam
#             else:
#                 used_cam = nao.bottom_cam
#
#             draw_str(bottom_img,(10,170),"point O in cam:")
#             tr2 = Tc_o.copy()
#             pc = tr2[:3,3]
#             p2d = used_cam._get_p2d_from_p3d(pc) # 相机坐标系下
#             draw_str(bottom_img,(10,190),str(p2d))
#
#             draw_str(bottom_img,(10,210),"point O in %s_cam:" % position)
#             p3d = Tb_o[:3,3] # 原点在机器人坐标系下的位置
#             p2d = used_cam.get_p2d_from_p3d(p3d)
#             draw_str(bottom_img,(10,230),str(p2d))
#
#         # 绘制左臂末端点
#         p3d = nao._getTf("LArm")[:3,3]
#         print("LArm:", p3d)
#         print("LArm Angles:", nao._getAngles("LArm"))
#         print("Head Angles:", nao._getAngles("Head"))
#         p2d = nao.top_cam.get_p2d_from_p3d(p3d.copy())
#         print("LArm top:", p2d)
#         plot_point(top_img,p2d,(255,255,0))
#         p2d = nao.bottom_cam.get_p2d_from_p3d(p3d)
#         print("LArm bottom:", p2d)
#         plot_point(bottom_img,p2d,(255,255,0))
#
#         # 显示图像
#         cv2.imshow('top', top_img)
#         cv2.imshow('bottom', bottom_img)
#         setAngle = lambda anglse: motionProxy.setAngles('Head', anglse, 0.5)
#         ch = cv2.waitKey(100) & 0xFF
#
#         command = {
#             "q": lambda step:setAngle([HeadYaw+step,HeadPitch]),
#             "a": lambda step:setAngle([HeadYaw-step,HeadPitch]),
#             "z": lambda step:setAngle([0,HeadPitch]),
#             "w": lambda step:setAngle([HeadYaw,HeadPitch+step]),
#             "s": lambda step:setAngle([HeadYaw,HeadPitch-step]),
#             "x": lambda step:setAngle([HeadYaw,0]),
#             "i": lambda step:autonomousLifeProxy.setState("disabled"),
#             "o": lambda step:motionProxy.wakeUp(),
#             "p": lambda step:motionProxy.rest(),
#             "j": lambda step:postureProxy.goToPosture("Stand", 1.0),
#             "k": lambda step:postureProxy.goToPosture("StandInit", 1.0),
#             "l": lambda step:postureProxy.goToPosture("Crouch", 1.0),
#         }
#         if ch == 27:
#             break
#         elif chr(ch) in command:
#             func = command[chr(ch)]
#             func(step)
#         elif ch == ord('e'): # 设定各角的极限
#             maxHeadYaw = HeadYaw
#         elif ch == ord('d'):
#             minHeadYaw = HeadYaw
#         elif ch == ord('r'):
#             maxHeadPitch = HeadPitch
#         elif ch == ord('f'):
#             minHeadPitch = HeadPitch
#         elif ch == ord("c"): # 保存当前
#             nao.states.append(sql.State("cam",nao,w,h,unit))
#             ss.add(nao)
#             ss.commit()
#             ss = sql.Session()
#             nao = get_nao()
#         elif ch == ord('v'): # 按照极限保存
#             ss = sql.Session()
#             for headYaw in np.linspace(minHeadYaw,maxHeadYaw,10):
#                 for headPitch in np.linspace(minHeadPitch,maxHeadPitch,10):
#                     print(headYaw,headPitch)
#                     setAngle([headYaw,headPitch])
#                     time.sleep(2)
#                     nao.states.append(sql.State("cam",nao,w,h,unit))
#                     ss.add(nao)
#                 ss.commit()
#             nao = get_nao()